package sink;

import com.asap.demo.utils.DateUtil;
import com.asap.demo.utils.Utils;
import com.asap.demo.utils.es.ESSinkUtil;
import com.asap.rule.StandardEvent;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
import org.apache.flink.streaming.api.TimeCharacteristic;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.test.util.MiniClusterWithClientResource;
import org.apache.http.HttpHost;
import org.junit.ClassRule;
import org.junit.Test;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.text.SimpleDateFormat;
import java.util.Date;
import java.util.List;
import java.util.Properties;

/**
 * @author wangbh
 * @Description: 入口
 * @date 2021/11/25 10:05
 */
public class FlinkConnectorsEs {
    private static final Logger logger = LoggerFactory.getLogger(FlinkConnectorsEs.class);

    @ClassRule
    public static MiniClusterWithClientResource flinkCluster =
            new MiniClusterWithClientResource(
                    new MiniClusterResourceConfiguration.Builder()
                            .setNumberSlotsPerTaskManager(3)
                            .setNumberTaskManagers(2)
                            .build());

    @Test
    public void test() throws Exception {
        try {
            final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

            env.setParallelism(1);
            env.enableCheckpointing(5000);  //检查点 每5000ms
            env.setStreamTimeCharacteristic(TimeCharacteristic.EventTime);

            Properties browseProperties = new Properties();
            browseProperties.put("bootstrap.servers", "192.168.1.25:9093");
            browseProperties.put("group.id", "temporal");
            browseProperties.put("auto.offset.reset", "latest");

            List<HttpHost> esAddresses = ESSinkUtil.getEsAddresses("http://192.168.1.239:9200");
            int bulkSize = 40;
            int sinkParallelism = 1;
            //

            DataStream<String> dataStream1 = env.fromElements("{\"\":\"1\",\"DST_MAC\":\"00:90:0b:48:87:51\",\"RAW_MSG\":\"<158>LEEF:1.0|Asiainfo security|TDA|3.83.1023|SECURITY_RISK_DETECTION|devTimeFormat=MMM dd yyyy HH:mm:ss z\\tptype=IDS\\tdvc=10.5.252.65\\tdeviceMacAddress=24:6E:96:17:C0:F4\\tdvchost=localhost\\tdeviceGUID=356D8743C9B3-47D1BC78-949C-22AD-A951\\tdevTime=Apr 19 2018 17:45:10 GMT+08:00\\tsev=8\\tprotoGroup=HTTP\\tproto=HTTP\\tvLANId=4095\\tdeviceDirection=1\\tdhost=125.66.84.35\\tdst=125.66.84.35\\tdstPort=80\\tdstMAC=00:1c:54:ff:08:02\\tshost=10.5.252.66\\tsrc=10.5.252.66\\tsrcPort=18384\\tsrcMAC=00:90:0b:48:87:51\\tmalType=OTHERS\\tfileType=-65536\\tfsize=0\\truleId=2348\\tmsg=CVE-2017-5638 - APACHE STRUTS EXPLOIT - HTTP (Request)\\tdeviceRiskConfidenceLevel=1\\tbotCommand=1071\\tbotUrl=iswin?{'cmd.exe','/c',#cmd}:{'/bin/bash','-c',#cmd})).(#p=new java.lang.ProcessBuilder(#cmds)).(#p.redirectErrorStream(true)).(#process=#p.start()).(#ros=(@org.apache.struts2.ServletActionContext@getResponse().getOutputStream())).(@org.apache.commons.io.I\\turl=http://jintang.ljcd.gov.cn/show-232-5192-1-14.html\\trequestClientApplication=Mozilla/5.0\\tpComp=NCIE\\triskType=1\\tsrcGroup=缺省\\tsrcZone=1\\tdstZone=0\\tdetectionType=1\\tact=not blocked\\tthreatType=3\\tinterestedIp=10.5.252.66\\tpeerIp=125.66.84.35\\thostName=jintang.ljcd.gov.cn\\tcnt=1\\taggregatedCnt=1\\tevtCat=Exploit\\tevtSubCat=Application\\taptRelated=0\\tpAttackPhase=Point of Entry\",\"SRC_PROVINCE\":\"四川省\",\"FILE_NAME\":\"\",\"SRC_LONGITUDE\":\"106.095100\",\"RISK_LEVEL_DESC\":\"高危\",\"DST_PROVINCE\":\"\",\"SNOW_ID\":\"1111977450112573447\",\"LOG_TYPE\":\"threat_log\",\"RISK_LEVEL\":\"3\",\"PARSED_SUCCESS\":\"true\",\"MALWARE_TYPE\":\"OTHERS\",\"RULE_ID\":\"2348\",\"EVENT_ONE_TYPE\":\"-1\",\"PROTOCOL\":\"HTTP\",\"DEVICE_IP\":\"127.0.0.1\",\"VULN_ID\":\"CVE-2017-5638\",\"pbadd\":\"pb\",\"MALDIUM_THREAT\":\"1\",\"LOG_SUB_TYPE\":\"SECURITY_RISK_DETECTION\",\"TAGS\":\"网络攻击,外连威胁\",\"MALDIUM_HASH_ID\":\"62a636520d92fa5efdb3a6ca3ab54043\",\"THREAT_TYPE\":\"3\",\"DOMAIN\":\"jintang.ljcd.gov.cn\",\"MAIL_SUBJECT\":\"\",\"DIRECTION_DESC\":\"外连\",\"MALDIUM_EN_ID\":\"BOZJYM4MbDlK7exQQELhQF59bQXRw6ItFydxRp6k6ErQr2ep9YmpJpMh8St2bTjHrNQix2scS00OusroEEb+QjmeCI3mFUH9liKqKXr9ZAaC3re3Temz77/aM5RwUYsYm6iuV85H+izVzN1z9Is=\",\"FLAG\":\"0\",\"MALWARE_NAME\":\"\",\"INFECTION_FILE\":\"\",\"URL\":\"http://jintang.ljcd.gov.cn/show-232-5192-1-14.html\",\"SRC_POST\":\"0\",\"MESSAGE\":\"10.5.252.66 对 125.66.84.35 发起'高危'的'CVE-2017-5638 - APACHE STRUTS EXPLOIT - HTTP (Request)'威胁,对其处理动作为'允许',处置结果为'成功'\",\"CREATE_TIME_TYPE\":\"2\",\"SRC_MAC\":\"00:1c:54:ff:08:02\",\"DST_HOST_NAME\":\"10.5.252.66\",\"collectTime\":\"2021-10-08 15:18:53.840\",\"DST_LATITUDE\":\"\",\"EVENT_NAME\":\"CVE-2017-5638 - APACHE STRUTS EXPLOIT - HTTP (Request)\",\"SRC_PORT\":\"80\",\"SRC_CITY\":\"南充市\",\"FILE_PATH\":\"\",\"filter_s\":\"0\",\"EVENT_TWO_TYPE_DESC\":\"漏洞攻击\",\"MAIL_SENDER\":\"\",\"recordTime\":\"2021-10-08 15:19:17.716\",\"TENANT_ID\":\"-1\",\"FILE_TYPE\":\"-65536\",\"TAG\":\"{\\\"ATTACK_STAGE\\\":\\\"入侵\\\",\\\"DIRECTION\\\":\\\"外连\\\",\\\"EVENT_THREE_TYPE\\\":\\\"远程命令/代码执行攻击\\\"}\",\"collectionName_s\":\"maxs_standard\",\"THREAT_TYPE_DESC\":\"漏洞利用\",\"EVENT_THREE_TYPE_DESC\":\"远程命令/代码执行攻击\",\"INTEL_NAME_ALL\":\"CVE-2017-5638,\",\"uuid\":\"896e2254-9d92-4ef0-a673-a083f0d64c001633677533840\",\"quality_rule_s\":\"0\",\"ACTION_DESC\":\"允许\",\"DST_LONGITUDE\":\"\",\"FILE_HASH\":\"\",\"FUNCTION_FLAG\":\"CUSTOM_VALUE1:strAdd(THREAT_TYPE,MALWARE_SUB_TYPE);CUSTOM_VALUE2:strAdd(THREAT_TYPE,MALWARE_TYPE);CUSTOM_VALUE3:strAtoA(RULE_ID)\",\"SRC_HOST_NAME\":\"125.66.84.35\",\"MALDIUM_TYPE_DESC\":\"maxs_cveid\",\"CREATE_TIME\":\"2021-10-08 17:19:17.528\",\"INTEL_ID_ALL\":\"e7061058e193d161aeb9e296e0425b7e,\",\"DEVICE_TYPE\":\"ASIA_TDA\",\"DST_COUNTRY\":\"\",\"NAME\":\"CVE-2017-5638\",\"MAIL_RECIPIENT\":\"\",\"CREATE_TIME_s\":\"Apr 19 2018 17:45:10\",\"MALWARE_SUB_TYPE\":\"Application\",\"MALD_COLLECT_DATE\":\"2021-10-08\",\"CREATE_TIME2\":\"2018-04-19 17:45:10.000\",\"MAIL_TYPE\":\"99\",\"EVENT_TWO_TYPE\":\"10000\",\"DEVICE_PARENT_TYPE\":\"TDA\",\"DST_IP\":\"10.5.252.66\",\"CLIENT_NAME\":\"localhost\",\"SRC_IP\":\"125.66.84.35\",\"CUSTOM_VALUE1\":\"3Application\",\"DEVICE_NAME\":\"深度威胁发现设备\",\"CUSTOM_VALUE2\":\"3OTHERS\",\"EVENT_THREE_TYPE\":\"10015\",\"SRC_LATITUDE\":\"30.797098\",\"MALDIUM_TYPE\":\"通用软硬件漏洞\",\"collectionName\":\"maxs_standard\",\"CUSTOM_VALUE9\":\"ASIA_TDA\",\"TAG1\":\"网络攻击\",\"INTEL_TYPE_ALL\":\"3,\",\"STRATEGY_ID\":\"10001\",\"CUSTOM_VALUE3\":\"_2348_\",\"DST_POST\":\"1\",\"TAG5\":\"外连威胁\",\"dpRuleId_s\":1101,\"DST_CITY\":\"\",\"DIRECTION\":\"1\",\"INTEL_ID\":\"e7061058e193d161aeb9e296e0425b7e\",\"MALDIUM_SOURCE\":\"\",\"INTEL_TYPE\":\"3\",\"ATTACK_STAGE\":\"入侵\",\"ACTION\":\"1\",\"equIP\":\"null\",\"MALDIUM_SEVERITY\":\"中\",\"DST_PORT\":\"18384\",\"SRC_COUNTRY\":\"中国\",\"CREATE_TIME2_s\":\"Apr 19 2018 17:45:10 GMT+08:00\",\"RESULT\":\"成功\",\"\":\"\"}");

            DataStream<StandardEvent> dataStream =dataStream1
                    .map(new MapFunction<String, StandardEvent>() {
                        private SimpleDateFormat simpleDate = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss.SSS");
                        private String indexCreateMode="DAY";
                        @Override
                        public StandardEvent map(String value) throws Exception {
                            StandardEvent standardEvent = StandardEvent.parse(value);
                            if(standardEvent.getAllFields().isEmpty()){
                                return null;
                            }
                            String indexName = null;
                            int count=0;
                            String collectTime = DateUtil.dateInfo();
                            if (standardEvent.getAllFields().containsKey("CREATE_TIME")) {
                                collectTime = (String) standardEvent.getField("CREATE_TIME");
                            }
                            standardEvent.getAllFields().put("collectTime", collectTime);
                            standardEvent.getAllFields().put("recordTime", simpleDate.format(new Date()));
                            standardEvent.getAllFields().put("rawMsg",value);
                            String collection = "maxs_standard";
                            if (standardEvent.getAllFields().containsKey("collectionName")) {
                                collection = (String) standardEvent.getField("collectionName");
                            }

                            Date date= DateUtil.parseDate(collectTime,DateUtil.DATE_FORMAT_FULL);
                            String time = DateUtil.parseDateToString(date,DateUtil.DATE_FORMAT);
                            if(logger.isDebugEnabled()){
                                logger.debug(Utils.join("","collectTime:",collectTime,"time:",time));

                            }

                            if("DAY".equals(indexCreateMode)){
                                indexName = Utils.join("_",collection,time.substring(0,8));
                            } else if("MONTH".equals(indexCreateMode)){
                                indexName = Utils.join("_",collection,time.substring(0,6));
                            } else if("HOUR".equals(indexCreateMode)){
                                indexName = Utils.join("_",collection,time.substring(0,10));
                            } else {
                                indexName = collection;
                            }
                            if (indexName.startsWith("-") || indexName.startsWith("_") || indexName.startsWith("+"))
                            {
                                count++;
                                if (count < 2)
                                {
                                    logger.error("",null,Utils.join("","当前设置的索引名称 ",
                                            collection," 初始化后名称 ",indexName," 不合法，其初始化索引名称变更为 unknown "));
                                }
                                indexName = "unknown";
                            }
                            standardEvent.getAllFields().put("indexName", indexName);

                            return standardEvent;
                        }
                    });

//            ElasticsearchSink.Builder<StandardEvent> esSinkBuilder = new ElasticsearchSink.Builder<>(esAddresses, new DealElasticsearchSinkFunction());
//            esSinkBuilder.setBulkFlushMaxActions(bulkSize);
//            esSinkBuilder.setFailureHandler(new RetryRequestFailureHandler());
//            //todo:xpack security
//            dataStream.addSink(esSinkBuilder.build()).setParallelism(sinkParallelism);
//            env.execute("flink connectors es");
        } catch (Exception e) {
            System.out.println(e);
        }
    }
}
